from dataclasses import dataclass, field
from typing import Dict, List
from .common import Simulator, CostCalculator

@dataclass
class OutputLayerSimulator(Simulator):
    def __init__(self, mp: Dict[str, object]):
        super().__init__()
        self.mp = mp
        self.hidden_size = mp['input_shape'][0][2]
        self.batch_size = mp['input_shape'][0][1]
        self.seq_len = mp['input_shape'][0][0]
        self.weight_hidden_size = mp['weight_shape'][0]
    
    def get_calculate_workload(self, tp_size: int) -> List[float]:
        calculator = [0,0,0,0,0]
        per_partition_output_size = self.weight_hidden_size // tp_size
        calculator[0] = self.batch_size * self.seq_len * ( self.hidden_size - 1 ) * per_partition_output_size
        calculator[2] = self.batch_size * self.seq_len * self.hidden_size * per_partition_output_size
        return calculator
    
    # gather 列切后的词表维度张量
    def get_communication_workload(self, tp_size) -> int:
        if tp_size == 1:
            return 0
        # 只用*一个4字节，输出层
        return self.batch_size * self.seq_len * self.weight_hidden_size * 4 / tp_size / 1024 / 1024 / 1024
    
    def get_dp_communication_workload(self, tp_size, dp_size):
        return 0
    
# 所有模型OutpurLayer的基类
@dataclass
class OutputLayer(CostCalculator):
    # 共有 列切
    output_layer: OutputLayerSimulator = None

    def statistic_single_layer_memory_workload(self, tp_size):
        # 输出层 与 embedding共享权重
        return 0, self.output_layer.mp['active_memory_allocated'] / 1024 / 1024 / 1024 / tp_size